package com.itmake.sample.limit.config;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.script.DefaultRedisScript;

/**
 * redis 限流脚本
 *
 * @author Amjacks
 * @since 2023/08/10
 */
@Configuration
public class RedisLimitScript {

    @Bean
    public DefaultRedisScript<Long> limitScript() {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptText(
                // 限流脚本
                "local key = KEYS[1]\n" +
                        "local count = tonumber(ARGV[1])\n" +
                        "local time = tonumber(ARGV[2])\n" +
                        "local current = redis.call('get', key);\n" +
                        "if current and tonumber(current) > count then\n" +
                        "    return tonumber(current);\n" +
                        "end\n" +
                        "current = redis.call('incr', key)\n" +
                        "if tonumber(current) == 1 then\n" +
                        "    redis.call('expire', key, time)\n" +
                        "end\n" +
                        "return tonumber(current);"
        );
        redisScript.setResultType(Long.class);
        return redisScript;
    }
}
